#!/anaconda3/envs/FEALPy/bin/python3.8
# -*- coding: utf-8 -*-
"""
File: qwen_client_new.py
Author: Bryan SHEN
E-mail: m18801919240_3@163.com
Site: Shanghai, China
Time: 2025/9/16
Description: 
"""

"""
qwen_client.py

一个兼容通义千问（DashScope）OpenAI-API 模式的轻量封装。
- 默认把 base_url 指向 DashScope 的 OpenAI 兼容端点（国内站）。
- 也可通过环境变量或参数自定义 base_url / api_key / model。
- 支持「场景」(summary / creative / default)、流式/非流式、超时与简单重试。
- 自带命令行 & 最小可运行测试样例。

快速使用：
  export DASHSCOPE_API_KEY="sk-xxxx"  # 或在参数/代码里传入
  # 内地（默认）：
  export OPENAI_BASE_URL="https://dashscope.aliyuncs.com/compatible-mode/v1"
  # 国际站（如你的 key 为国际站）：
  # export OPENAI_BASE_URL="https://dashscope-intl.aliyuncs.com/compatible-mode/v1"

  python qwen_client.py --prompt "请用中文总结《百年孤独》" --scene summary
  python qwen_client.py --prompt "写一首七言绝句" --scene creative --stream
"""


import os
import sys
import time
import json
import argparse
from typing import List, Dict, Optional, Iterable, Union

from tenacity import retry, stop_after_attempt, wait_exponential
import httpx


# -----------------------------
# 默认配置（可按需修改）
# -----------------------------
DEFAULT_CN_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
DEFAULT_INTL_BASE_URL = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
DEFAULT_MODEL = os.getenv("QWEN_MODEL", "qwen-plus")
# 内置默认 API Key（如未提供环境变量与参数，则使用此值）
DEFAULT_API_KEY = os.getenv("DASHSCOPE_API_KEY") or os.getenv("OPENAI_API_KEY") or "sk-0342f3165fc04dc5b306d4acb1d898df"

SCENE_SYSTEM_PROMPTS: Dict[str, str] = {
    "default": "You are a helpful assistant.",
    "summary": "你是一个精炼、客观的中文摘要助手。要求：忠实原文、结构清晰、要点明确、避免主观评判。",
    "creative": "你是一位中文创作助手，文风可适度抒情与讲究节奏，但需保持清晰与可读性。",
    "factual": "You are a factual and concise assistant. Answer precisely based on provided data.",
}


def _resolve_base_url(user_base_url: Optional[str]) -> str:
    """
    决定实际使用的 base_url：
    优先级：入参 > OPENAI_BASE_URL > DASHSCOPE_BASE_URL > DEFAULT_CN_BASE_URL
    """
    if user_base_url:
        return user_base_url
    env_url = os.getenv("OPENAI_BASE_URL") or os.getenv("DASHSCOPE_BASE_URL")
    return env_url or DEFAULT_CN_BASE_URL


def _resolve_api_key(user_key: Optional[str]) -> str:
    """
    决定实际使用的 api_key：
    优先级：入参 > DASHSCOPE_API_KEY > OPENAI_API_KEY
    """
    key = user_key or DEFAULT_API_KEY
    if not key:
        raise RuntimeError(
            "缺少 API Key。请通过参数 --api-key 或环境变量 DASHSCOPE_API_KEY / OPENAI_API_KEY 提供。"
        )
    return key


def _build_messages(
    prompt: Union[str, List[Dict[str, str]]],
    scene: Optional[str] = None,
    system_prompt: Optional[str] = None,
) -> List[Dict[str, str]]:
    """
    将 prompt 与场景转换为 OpenAI Chat Completions 的 messages。
    - 如果 prompt 是 str：拼接 system + user。
    - 如果 prompt 是 list[message]：前面附加 system。
    """
    scene_key = (scene or "default").lower()
    sys_prompt = system_prompt or SCENE_SYSTEM_PROMPTS.get(scene_key, SCENE_SYSTEM_PROMPTS["default"])

    messages = [{"role": "system", "content": sys_prompt}]
    if isinstance(prompt, str):
        messages.append({"role": "user", "content": prompt})
    elif isinstance(prompt, list):
        messages.extend(prompt)
    else:
        raise TypeError("prompt 必须是 str 或 list[dict(role, content)]")
    return messages


class QwenClient:
    """
    通义千问（DashScope）OpenAI-API 兼容模式客户端封装。
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: str = DEFAULT_MODEL,
        timeout: float = 60.0,
        max_retries: int = 2,
    ) -> None:
        self.api_key = _resolve_api_key(api_key)
        self.base_url = _resolve_base_url(base_url)
        self.model = model
        self.timeout = timeout
        self.max_retries = max_retries

        # 清理常见代理环境变量，避免 httpx 注入 proxies 触发不兼容
        for var in ("HTTP_PROXY","HTTPS_PROXY","http_proxy","https_proxy","ALL_PROXY","all_proxy"):
            if var in os.environ:
                try:
                    del os.environ[var]
                except Exception:
                    pass

    @retry(wait=wait_exponential(multiplier=1, min=1, max=8), stop=stop_after_attempt(3), reraise=True)
    def _create_completion(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: Optional[int] = None,
        stream: bool = False,
    ):
        """通过 OpenAI 兼容 REST 端点发起请求（禁用环境代理）。"""
        url = f"{self.base_url}/chat/completions"
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": float(temperature),
            "top_p": float(top_p),
            "max_tokens": max_tokens,
            "stream": bool(stream),
        }
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        with httpx.Client(timeout=self.timeout, trust_env=False) as client:
            resp = client.post(url, headers=headers, json=payload)
            # 对于 401/403 给出更明确提示
            if resp.status_code == 401:
                raise RuntimeError("鉴权失败(401)。请检查 DASHSCOPE_API_KEY / OPENAI_BASE_URL 是否匹配。")
            if resp.status_code == 403:
                raise RuntimeError("权限不足(403)。请检查账户权限与模型可用性。")
            if resp.status_code in (429, 500, 502, 503, 504):
                # 触发 tenacity 重试
                raise RuntimeError(f"服务繁忙或错误: {resp.status_code}")
            resp.raise_for_status()
            return resp

    def ask(
        self,
        prompt: Union[str, List[Dict[str, str]]],
        scene: Optional[str] = None,
        system_prompt: Optional[str] = None,
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        return_usage: bool = False,
    ) -> Union[str, Iterable[str], Dict[str, Union[str, Dict]]]:
        """
        发起一次对话请求。
        - 非流式：返回完整 string（或包含 text + usage 的 dict）
        - 流式：返回可迭代的片段（yield str）

        return_usage=True 时，非流式返回：
        {
          "text": "...",
          "usage": {...},
        }
        """
        messages = _build_messages(prompt, scene=scene, system_prompt=system_prompt)

        try:
            if stream:
                # 简单流式：对非流式返回的完整文本逐字符回放
                resp = self._create_completion(
                    messages=messages,
                    temperature=temperature,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    stream=False,
                )
                data = resp.json()
                text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                def _gen():
                    for ch in text:
                        yield ch
                return _gen()

            # 非流式
            resp = self._create_completion(
                messages=messages,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                stream=False,
            )
            data = resp.json()
            text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
            if return_usage:
                usage = None
                try:
                    u = data.get("usage")
                    if isinstance(u, dict):
                        usage = {
                            "prompt_tokens": u.get("prompt_tokens"),
                            "completion_tokens": u.get("completion_tokens"),
                            "total_tokens": u.get("total_tokens"),
                        }
                except Exception:
                    usage = None
                return {"text": text, "usage": usage}
            return text
        except Exception:
            # 直接抛给上层，便于 UI 显示真实错误
            raise


# -----------------------------
# 测试 & 命令行
# -----------------------------
def _smoke_test_non_stream(qc: QwenClient) -> None:
    print("== 非流式摘要测试 ==")
    out = qc.ask("请用中文总结《百年孤独》，不超过120字，给出3个要点。", scene="summary", return_usage=True)
    print(json.dumps(out, ensure_ascii=False, indent=2))


def _smoke_test_stream(qc: QwenClient) -> None:
    print("\n== 流式创作测试 ==")
    gen = qc.ask("以“海风”“旧街”“白昼”写一段100字以内抒情短文。", scene="creative", stream=True)
    for chunk in gen:
        print(chunk, end="", flush=True)
    print("\n[流式结束]")


def main():

    parser = argparse.ArgumentParser(description="QwenClient 命令行示例")
    # parser.add_argument("--api-key", type=str, default=None, help="API Key（默认读环境变量）")
    parser.add_argument("--api-key", type=str, default="sk-0342f3165fc04dc5b306d4acb1d898df", help="API Key（默认读环境变量）")
    parser.add_argument("--base-url", type=str, default=None, help="DashScope 兼容端点（默认自动）")
    parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help=f"模型名（默认 {DEFAULT_MODEL}）")
    parser.add_argument("--prompt", type=str, default="请用中文总结《百年孤独》", help="用户输入")
    parser.add_argument("--scene", type=str, default="summary", choices=list(SCENE_SYSTEM_PROMPTS.keys()), help="场景（影响system提示词）")
    parser.add_argument("--system", type=str, default=None, help="自定义 system 提示词，优先于 scene")
    parser.add_argument("--temperature", type=float, default=0.7, help="采样温度")
    parser.add_argument("--top-p", type=float, default=0.9, help="核采样阈值")
    parser.add_argument("--max-tokens", type=int, default=None, help="最大输出 token 数")
    parser.add_argument("--stream", action="store_true", help="是否流式输出")
    parser.add_argument("--smoke", action="store_true", help="运行内置 smoke tests")
    args = parser.parse_args()

    qc = QwenClient(
        api_key=args.api_key,
        base_url=args.base_url,
        model=args.model,
        timeout=60.0,
        max_retries=2,
    )

    if args.smoke:
        _smoke_test_non_stream(qc)
        _smoke_test_stream(qc)
        return

    # 单次调用（命令行）
    if args.stream:
        for chunk in qc.ask(
            prompt=args.prompt,
            scene=args.scene,
            system_prompt=args.system,
            temperature=args.temperature,
            top_p=args.top_p,
            max_tokens=args.max_tokens,
            stream=True,
        ):
            print(chunk, end="", flush=True)
        print()
    else:
        out = qc.ask(
            prompt=args.prompt,
            scene=args.scene,
            system_prompt=args.system,
            temperature=args.temperature,
            top_p=args.top_p,
            max_tokens=args.max_tokens,
            stream=False,
            return_usage=True,
        )
        print(json.dumps(out, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        # 统一打印更友好的错误信息
        print(f"[ERROR] {e}", file=sys.stderr)
        sys.exit(1)

